import torch
import torch.nn as nn
import torch.nn.functional as F


class GraphConvolutionLayer(torch.nn.Module):
    def __init__(self, in_features, out_features):
        super(GraphConvolutionLayer, self).__init__()
        self.weight = torch.nn.Parameter(
            torch.FloatTensor(in_features, out_features))
        self.bias = torch.nn.Parameter(torch.FloatTensor(out_features))
        torch.nn.init.xavier_uniform_(self.weight)

    def forward(self, adjacency_matrix, input_features):
        support = input_features@self.weight
        output = adjacency_matrix@support + self.bias
        return output


class GlobalPoolingGNN(torch.nn.Module):
    def __init__(self, num_nodes, in_features, hidden_features, out_features):
        super(GlobalPoolingGNN, self).__init__()
        self.conv1 = GraphConvolutionLayer(in_features, hidden_features)
        self.conv2 = GraphConvolutionLayer(hidden_features, out_features)
        self.out_features = out_features
        self.num_nodes = num_nodes

    def forward(self, adjacency_matrix, input_features):
        x = F.relu(self.conv1(adjacency_matrix, input_features))
        x = F.relu(self.conv2(adjacency_matrix, x))

        # Global pooling (mean pooling in this example)
        x = torch.mean(x, dim=1, keepdim=True)
        x = x.view((-1, self.out_features))

        return x


class GNNClassifier(torch.nn.Module):
    def __init__(self, num_nodes, in_features, hidden_features, out_features, class_num):
        super(GNNClassifier, self).__init__()
        self.gnn = GlobalPoolingGNN(
            num_nodes, in_features, hidden_features, out_features)
        self.lr = nn.Linear(out_features, class_num)
        self.num_nodes = num_nodes

    def forward(self, adjacency_matrix, input_features):
        x = self.gnn(adjacency_matrix, input_features)
        x = self.lr(x)
        return x
